import numpy as np
from collections import Counter


class KNN:
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):
        self.X_train = np.array(X)
        self.y_train = np.array(y)

    def euclidean_distance(self, x1, x2):
        return np.sqrt(np.sum((x1 - x2) ** 2))

    def manhattan_distance(self, x1, x2):
        return np.sum(np.abs(x1 - x2))

    def _get_neighbors(self, x):
        distances = []
        for i, train_sample in enumerate(self.X_train):
            dist = self.euclidean_distance(x, train_sample)
            distances.append((i, dist))
        distances.sort(key=lambda x: x[1])
        return distances[:self.k]

    def predict(self, X):
        X = np.array(X)
        predictions = []

        for sample in X:
            neighbors = self._get_neighbors(sample)
            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
            most_common = Counter(neighbor_labels).most_common(1)
            predictions.append(most_common[0][0])

        return np.array(predictions)

    def predict_proba(self, X):
        X = np.array(X)
        probabilities = []
        unique_classes = np.unique(self.y_train)

        for sample in X:
            neighbors = self._get_neighbors(sample)
            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
            label_counts = Counter(neighbor_labels)
            proba = [label_counts.get(cls, 0) / self.k for cls in unique_classes]
            probabilities.append(proba)

        return np.array(probabilities)

    def score(self, X, y):
        predictions = self.predict(X)
        return np.mean(predictions == y)